package com.shanshui.test.gm;


import com.tencent.kona.crypto.KonaCryptoProvider;
import com.tencent.kona.pkix.KonaPKIXProvider;
import com.tencent.kona.ssl.KonaSSLProvider;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.security.Security;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Collection;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;

/**
 * @author mashanshui
 * @since 2025/11/14
 */
public class GMSSLContext {
    private TrustManager[] trustManagers = null;
    private SSLContext sslContext = null;

    static {
        Security.addProvider(new KonaCryptoProvider());
        Security.addProvider(new KonaPKIXProvider());
        Security.addProvider(new KonaSSLProvider());
    }

    public GMSSLContext(String certChain) {
        this(new ByteArrayInputStream(certChain.getBytes()));
    }

    public GMSSLContext(InputStream certChain) {
        try {
            KeyStore trustStore = createTrustStore(certChain);
            TrustManagerFactory tmf = TrustManagerFactory.getInstance("PKIX", "KonaSSL");
            tmf.init(trustStore);
            trustManagers = tmf.getTrustManagers();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        initContext();
    }

    public GMSSLContext(TrustManager[] trustManagers) {
        this.trustManagers = trustManagers;
        initContext();
    }

    private void initContext() {
        try {
            SSLContext context = SSLContext.getInstance("TLCPv1.1", "KonaSSL");
            context.init(null, trustManagers, new SecureRandom());
            sslContext = context;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }


    private static KeyStore createTrustStore(InputStream certIn)
            throws Exception {
        KeyStore trustStore = KeyStore.getInstance("PKCS12", "KonaPKIX");
        trustStore.load(null, null);
        CertificateFactory cf = CertificateFactory.getInstance("X.509", "KonaPKIX");
        Collection<? extends Certificate> certs = cf.generateCertificates(certIn);

        int index = 0;
        for (Certificate cert : certs) {
            if (cert instanceof X509Certificate) {
                // 将每个CA证书添加到信任库
                trustStore.setCertificateEntry("ca-alias-" + index++, cert);
            }
        }
        return trustStore;
    }

    public TrustManager[] getTrustManagers() {
        return trustManagers;
    }

    public SSLContext getSslContext() {
        return sslContext;
    }

    public SSLSocketFactory getSslSocketFactory() {
        return sslContext.getSocketFactory();
    }
}
